"""
WebSocket实时推送服务
==================

基于FastAPI WebSocket的实时市场数据推送服务
提供连接管理、订阅管理、数据广播等功能
"""

import asyncio
import json
import logging
import time
from typing import Dict, Set, List, Any, Optional, Callable
from datetime import datetime, timezone
from dataclasses import dataclass, asdict
from enum import Enum

from fastapi import WebSocket, WebSocketDisconnect
from pydantic import BaseModel, Field

logger = logging.getLogger(__name__)


class MessageType(str, Enum):
    """WebSocket消息类型"""
    CONNECT = "connect"
    DISCONNECT = "disconnect"
    SUBSCRIBE = "subscribe"
    UNSUBSCRIBE = "unsubscribe"
    DATA = "data"
    ERROR = "error"
    PING = "ping"
    PONG = "pong"
    HEARTBEAT = "heartbeat"


class WebSocketMessage(BaseModel):
    """WebSocket消息模型"""
    type: MessageType
    data: Dict[str, Any] = Field(default_factory=dict)
    timestamp: float = Field(default_factory=time.time)
    request_id: Optional[str] = None


@dataclass
class ConnectionInfo:
    """连接信息"""
    websocket: WebSocket
    client_id: str
    user_id: Optional[str] = None
    subscriptions: Set[str] = None
    connected_at: datetime = None
    last_ping: float = None
    
    def __post_init__(self):
        if self.subscriptions is None:
            self.subscriptions = set()
        if self.connected_at is None:
            self.connected_at = datetime.now(timezone.utc)
        if self.last_ping is None:
            self.last_ping = time.time()


class WebSocketService:
    """WebSocket连接管理服务"""
    
    def __init__(self, max_connections: int = 1000):
        self.max_connections = max_connections
        self.connections: Dict[str, ConnectionInfo] = {}
        self.subscriptions: Dict[str, Set[str]] = {}  # symbol -> client_ids
        self.data_callbacks: List[Callable] = []
        self.metrics = WebSocketMetrics()
        
        # 心跳检测任务（延迟初始化）
        self._heartbeat_task = None
        self._initialized = False
    
    async def connect(self, websocket: WebSocket, client_id: str, user_id: Optional[str] = None) -> bool:
        """
        建立WebSocket连接
        
        Args:
            websocket: WebSocket连接对象
            client_id: 客户端唯一标识
            user_id: 用户ID（可选）
            
        Returns:
            bool: 连接是否成功
        """
        # 首次连接时初始化心跳任务
        if not self._initialized:
            self._start_heartbeat()
            self._initialized = True
        
        try:
            # 检查连接数限制
            if len(self.connections) >= self.max_connections:
                logger.warning(f"WebSocket连接数已达上限: {self.max_connections}")
                await websocket.close(code=1013, reason="Server overloaded")
                return False
            
            # 检查客户端是否已连接
            if client_id in self.connections:
                logger.warning(f"客户端 {client_id} 已存在连接，关闭旧连接")
                await self.disconnect(client_id)
            
            # 接受WebSocket连接
            await websocket.accept()
            
            # 创建连接信息
            connection = ConnectionInfo(
                websocket=websocket,
                client_id=client_id,
                user_id=user_id
            )
            
            # 保存连接
            self.connections[client_id] = connection
            self.metrics.record_connection()
            
            # 发送连接成功消息
            await self._send_message(websocket, MessageType.CONNECT, {
                "client_id": client_id,
                "message": "WebSocket连接已建立",
                "server_time": datetime.now(timezone.utc).isoformat()
            })
            
            logger.info(f"WebSocket连接已建立: {client_id}")
            return True
            
        except Exception as e:
            logger.error(f"WebSocket连接失败: {e}", exc_info=True)
            self.metrics.record_error("connection_failed")
            return False
    
    async def disconnect(self, client_id: str) -> bool:
        """
        断开WebSocket连接
        
        Args:
            client_id: 客户端ID
            
        Returns:
            bool: 断开是否成功
        """
        try:
            if client_id not in self.connections:
                return False
            
            connection = self.connections[client_id]
            
            # 取消所有订阅
            for symbol in connection.subscriptions.copy():
                await self.unsubscribe(client_id, symbol)
            
            # 关闭WebSocket连接
            try:
                await connection.websocket.close()
            except Exception:
                pass  # 连接可能已经关闭
            
            # 移除连接
            del self.connections[client_id]
            self.metrics.record_disconnection()
            
            logger.info(f"WebSocket连接已断开: {client_id}")
            return True
            
        except Exception as e:
            logger.error(f"WebSocket断开失败: {e}", exc_info=True)
            self.metrics.record_error("disconnection_failed")
            return False
    
    async def subscribe(self, client_id: str, symbols: List[str]) -> bool:
        """
        订阅数据
        
        Args:
            client_id: 客户端ID
            symbols: 要订阅的股票代码列表
            
        Returns:
            bool: 订阅是否成功
        """
        try:
            if client_id not in self.connections:
                return False
            
            connection = self.connections[client_id]
            
            for symbol in symbols:
                # 添加到客户端订阅列表
                connection.subscriptions.add(symbol)
                
                # 添加到全局订阅映射
                if symbol not in self.subscriptions:
                    self.subscriptions[symbol] = set()
                self.subscriptions[symbol].add(client_id)
            
            # 发送订阅成功消息
            await self._send_to_client(client_id, MessageType.SUBSCRIBE, {
                "symbols": symbols,
                "message": f"已订阅 {len(symbols)} 个股票代码"
            })
            
            logger.info(f"客户端 {client_id} 订阅了 {symbols}")
            return True
            
        except Exception as e:
            logger.error(f"订阅失败: {e}", exc_info=True)
            self.metrics.record_error("subscribe_failed")
            return False
    
    async def unsubscribe(self, client_id: str, symbols: List[str]) -> bool:
        """
        取消订阅
        
        Args:
            client_id: 客户端ID
            symbols: 要取消订阅的股票代码列表
            
        Returns:
            bool: 取消订阅是否成功
        """
        try:
            if client_id not in self.connections:
                return False
            
            connection = self.connections[client_id]
            
            for symbol in symbols:
                # 从客户端订阅列表移除
                connection.subscriptions.discard(symbol)
                
                # 从全局订阅映射移除
                if symbol in self.subscriptions:
                    self.subscriptions[symbol].discard(client_id)
                    if not self.subscriptions[symbol]:
                        del self.subscriptions[symbol]
            
            # 发送取消订阅成功消息
            await self._send_to_client(client_id, MessageType.UNSUBSCRIBE, {
                "symbols": symbols,
                "message": f"已取消订阅 {len(symbols)} 个股票代码"
            })
            
            logger.info(f"客户端 {client_id} 取消订阅了 {symbols}")
            return True
            
        except Exception as e:
            logger.error(f"取消订阅失败: {e}", exc_info=True)
            self.metrics.record_error("unsubscribe_failed")
            return False
    
    async def broadcast_data(self, symbol: str, data: Dict[str, Any]) -> int:
        """
        广播数据到订阅者
        
        Args:
            symbol: 股票代码
            data: 要广播的数据
            
        Returns:
            int: 成功推送的客户端数量
        """
        try:
            if symbol not in self.subscriptions:
                return 0
            
            # 准备消息
            message = WebSocketMessage(
                type=MessageType.DATA,
                data={
                    "symbol": symbol,
                    "data": data,
                    "timestamp": time.time()
                }
            )
            
            # 推送给所有订阅者
            success_count = 0
            failed_clients = []
            
            for client_id in self.subscriptions[symbol].copy():
                try:
                    if await self._send_to_client(client_id, message.type, message.data):
                        success_count += 1
                    else:
                        failed_clients.append(client_id)
                except Exception as e:
                    logger.warning(f"推送数据到客户端 {client_id} 失败: {e}")
                    failed_clients.append(client_id)
            
            # 清理失败的连接
            for client_id in failed_clients:
                await self.disconnect(client_id)
            
            # 记录指标
            self.metrics.record_message(len(json.dumps(data)))
            
            logger.debug(f"向 {success_count} 个客户端推送了 {symbol} 的数据")
            return success_count
            
        except Exception as e:
            logger.error(f"广播数据失败: {e}", exc_info=True)
            self.metrics.record_error("broadcast_failed")
            return 0
    
    async def broadcast_to_all(self, data: Dict[str, Any]) -> int:
        """
        广播数据到所有连接
        
        Args:
            data: 要广播的数据
            
        Returns:
            int: 成功推送的客户端数量
        """
        try:
            message = WebSocketMessage(
                type=MessageType.DATA,
                data=data
            )
            
            success_count = 0
            failed_clients = []
            
            for client_id in list(self.connections.keys()):
                try:
                    if await self._send_to_client(client_id, message.type, message.data):
                        success_count += 1
                    else:
                        failed_clients.append(client_id)
                except Exception as e:
                    logger.warning(f"广播数据到客户端 {client_id} 失败: {e}")
                    failed_clients.append(client_id)
            
            # 清理失败的连接
            for client_id in failed_clients:
                await self.disconnect(client_id)
            
            self.metrics.record_message(len(json.dumps(data)))
            return success_count
            
        except Exception as e:
            logger.error(f"全局广播失败: {e}", exc_info=True)
            self.metrics.record_error("broadcast_all_failed")
            return 0
    
    async def handle_message(self, client_id: str, message: str) -> bool:
        """
        处理客户端消息
        
        Args:
            client_id: 客户端ID
            message: 消息内容
            
        Returns:
            bool: 处理是否成功
        """
        try:
            data = json.loads(message)
            msg_type = data.get("type")
            msg_data = data.get("data", {})
            
            if msg_type == MessageType.SUBSCRIBE:
                symbols = msg_data.get("symbols", [])
                return await self.subscribe(client_id, symbols)
            
            elif msg_type == MessageType.UNSUBSCRIBE:
                symbols = msg_data.get("symbols", [])
                return await self.unsubscribe(client_id, symbols)
            
            elif msg_type == MessageType.PING:
                await self._send_to_client(client_id, MessageType.PONG, {
                    "timestamp": time.time()
                })
                return True
            
            elif msg_type == MessageType.HEARTBEAT:
                if client_id in self.connections:
                    self.connections[client_id].last_ping = time.time()
                return True
            
            else:
                logger.warning(f"未知消息类型: {msg_type}")
                await self._send_to_client(client_id, MessageType.ERROR, {
                    "message": f"未知消息类型: {msg_type}"
                })
                return False
                
        except json.JSONDecodeError:
            logger.warning(f"无效的JSON消息: {message}")
            await self._send_to_client(client_id, MessageType.ERROR, {
                "message": "无效的JSON格式"
            })
            return False
        except Exception as e:
            logger.error(f"处理消息失败: {e}", exc_info=True)
            self.metrics.record_error("message_handle_failed")
            return False
    
    async def _send_to_client(self, client_id: str, msg_type: MessageType, data: Dict[str, Any]) -> bool:
        """发送消息到指定客户端"""
        try:
            if client_id not in self.connections:
                return False
            
            connection = self.connections[client_id]
            return await self._send_message(connection.websocket, msg_type, data)
            
        except Exception as e:
            logger.warning(f"发送消息到客户端 {client_id} 失败: {e}")
            return False
    
    async def _send_message(self, websocket: WebSocket, msg_type: MessageType, data: Dict[str, Any]) -> bool:
        """发送WebSocket消息"""
        try:
            message = WebSocketMessage(type=msg_type, data=data)
            await websocket.send_text(message.model_dump_json())
            return True
        except Exception as e:
            logger.warning(f"发送WebSocket消息失败: {e}")
            return False
    
    def _start_heartbeat(self):
        """启动心跳检测任务"""
        if self._heartbeat_task is not None:
            return  # 已经启动
            
        async def heartbeat_loop():
            while True:
                try:
                    await asyncio.sleep(30)  # 每30秒检查一次
                    await self._check_heartbeat()
                except Exception as e:
                    logger.error(f"心跳检测失败: {e}")
        
        try:
            self._heartbeat_task = asyncio.create_task(heartbeat_loop())
        except RuntimeError:
            # 如果没有运行的事件循环，稍后会在connect时启动
            logger.debug("心跳任务将在首次连接时启动")
    
    async def _check_heartbeat(self):
        """检查心跳，清理超时连接"""
        current_time = time.time()
        timeout_clients = []
        
        for client_id, connection in self.connections.items():
            if current_time - connection.last_ping > 60:  # 60秒超时
                timeout_clients.append(client_id)
        
        for client_id in timeout_clients:
            logger.info(f"客户端 {client_id} 心跳超时，断开连接")
            await self.disconnect(client_id)
    
    def get_connection_count(self) -> int:
        """获取当前连接数"""
        return len(self.connections)
    
    def get_subscription_count(self) -> int:
        """获取订阅数"""
        return sum(len(subs) for subs in self.subscriptions.values())
    
    def get_all_subscriptions(self) -> Set[str]:
        """
        获取所有已订阅的标的代码
        
        Returns:
            Set[str]: 所有订阅的标的代码集合
        """
        # 返回所有被订阅的股票代码（keys），而不是客户端ID（values）
        return set(self.subscriptions.keys())
    
    def get_metrics(self) -> Dict[str, Any]:
        """获取服务指标"""
        return self.metrics.get_metrics()


class WebSocketMetrics:
    """WebSocket监控指标"""
    
    def __init__(self):
        self.active_connections = 0
        self.total_connections = 0
        self.total_disconnections = 0
        self.total_messages_sent = 0
        self.total_bytes_sent = 0
        self.error_count = 0
        self.error_types = {}
        self.start_time = time.time()
    
    def record_connection(self):
        """记录连接建立"""
        self.active_connections += 1
        self.total_connections += 1
    
    def record_disconnection(self):
        """记录连接断开"""
        self.active_connections = max(0, self.active_connections - 1)
        self.total_disconnections += 1
    
    def record_message(self, size: int):
        """记录消息发送"""
        self.total_messages_sent += 1
        self.total_bytes_sent += size
    
    def record_error(self, error_type: str):
        """记录错误"""
        self.error_count += 1
        self.error_types[error_type] = self.error_types.get(error_type, 0) + 1
    
    def get_metrics(self) -> Dict[str, Any]:
        """获取指标数据"""
        uptime = time.time() - self.start_time
        return {
            "active_connections": self.active_connections,
            "total_connections": self.total_connections,
            "total_disconnections": self.total_disconnections,
            "total_messages_sent": self.total_messages_sent,
            "total_bytes_sent": self.total_bytes_sent,
            "error_count": self.error_count,
            "error_types": self.error_types,
            "uptime_seconds": uptime,
            "messages_per_second": self.total_messages_sent / uptime if uptime > 0 else 0,
            "bytes_per_second": self.total_bytes_sent / uptime if uptime > 0 else 0
        }


# 全局WebSocket服务实例
websocket_service = WebSocketService()
